#include "chess_board_distance.h"

using namespace std;
using namespace cv;

void ChessBoardDistance::getObjectPoints(std::vector<cv::Point3f>& objectPoints)
{
	const float squareLength = parameters["square_length"];
	objectPoints.resize(4);
	objectPoints[0] = { -squareLength / 2, squareLength / 2, 0 };
	objectPoints[1] = { squareLength / 2, squareLength / 2, 0 };
	objectPoints[2] = { squareLength / 2, -squareLength / 2, 0 };
	objectPoints[3] = { -squareLength / 2, -squareLength / 2, 0 };
}

void ChessBoardDistance::getImagePoints(const cv::Mat& image, std::vector<cv::Point2f>& imagePoints)
{
	Mat grayImage;
	if (image.channels() == 3) {
		cvtColor(image, grayImage, COLOR_BGR2GRAY);
	} else {
		grayImage = image.clone();
	}

	vector<Point2f> corners;
	bool found = findChessboardCorners(grayImage, { 3, 3 }, corners,
		CALIB_CB_ADAPTIVE_THRESH | CALIB_CB_FAST_CHECK | CALIB_CB_NORMALIZE_IMAGE);
	if (!found) {
		return;
	}

	cornerSubPix(grayImage, corners, Size(11, 11), Size(-1, -1),
		TermCriteria(TermCriteria::EPS | TermCriteria::MAX_ITER, 30, 0.001));

	imagePoints.resize(4);
	imagePoints[0] = corners[0];
	imagePoints[1] = corners[2];
	imagePoints[2] = corners[8];
	imagePoints[3] = corners[6];
}

double ChessBoardDistance::calculate(cv::Mat& image)
{
	Mat& cameraMatrix = cameraParameter.cameraMatrix;
	Mat& distCoeffs = cameraParameter.distCoeffs;
	const double squareLength = parameters["square_length"];
	const double scale = parameters["scale"];
	const double reference = parameters["reference"];

	if (cameraMatrix.empty() || distCoeffs.empty() || squareLength <= 0 || scale <= 0) {
		return -1;
	}

	vector<Point3f> objectPoints;
	getObjectPoints(objectPoints);

	vector<Point2f> imagePoints;
	getImagePoints(image, imagePoints);

	if (objectPoints.empty() || imagePoints.empty() || objectPoints.size() != imagePoints.size()) {
		return -1;
	}

	cv::Mat rvec, tvec;
	bool retval = solvePnP(objectPoints, imagePoints, cameraMatrix, distCoeffs,
		rvec, tvec, false, SOLVEPNP_IPPE_SQUARE);
	if (!retval) {
		return -1;
	}

	cv::Mat rotateMatrix;
	Rodrigues(rvec, rotateMatrix);
	cv::Mat transMatrix = -rotateMatrix.inv() * tvec;

	return transMatrix.at<double>(2) * scale + reference;
}


